from datetime import datetime
from typing import List, Optional
import json

from loguru import logger
from sqlalchemy import Column, String, DateTime, Integer, Text
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.declarative import declarative_base


Base = declarative_base()


class ChatSession(Base):
    """聊天会话表"""
    __tablename__ = 'chat_sessions'

    session_id = Column(String(64), primary_key=True, index=True)
    created_at = Column(DateTime, default=datetime.now(), nullable=False)
    last_activity = Column(DateTime, default=datetime.now(), nullable=False)
    message_count = Column(Integer, default=0, nullable=False)
    last_message = Column(Text, nullable=True)
    last_message_role = Column(String(20), nullable=True)
    chat_metadata = Column(Text, nullable=True)
    knowledge_base_ids = Column(Text, nullable=True)  # 新增：存储知识库ID列表的JSON字符串

    def get_knowledge_base_ids(self) -> List[str]:
        """获取知识库ID列表"""
        if not self.knowledge_base_ids:
            return []
        try:
            return json.loads(self.knowledge_base_ids)
        except (json.JSONDecodeError, TypeError):
            return []

    def set_knowledge_base_ids(self, kb_ids: List[str]):
        """设置知识库ID列表"""
        if kb_ids:
            self.knowledge_base_ids = json.dumps(kb_ids)
        else:
            self.knowledge_base_ids = None


class SessionManager:
    """会话管理器"""

    def __init__(self):
        self.engine = None

    async def _get_engine(self):
        """获取数据库引擎"""
        if self.engine is None:
            from ai_platform.config.resource import get_db_engine
            self.engine = get_db_engine()
        return self.engine

    async def init_tables(self):
        """初始化数据库表"""
        try:
            engine = await self._get_engine()
            async with engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
            logger.info("会话表初始化完成")
        except Exception as e:
            logger.exception(f"会话表初始化失败: {e}")
            raise

    async def create_session(self, session_id: str, knowledge_base_ids: Optional[List[str]] = None) -> ChatSession:
        """创建新会话"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                # 检查会话是否已存在
                existing = await session.get(ChatSession, session_id)
                if existing:
                    return existing
                # 创建新会话
                new_session = ChatSession(
                    session_id=session_id,
                    created_at=datetime.now(),
                    last_activity=datetime.now(),
                    message_count=0
                )
                
                # 设置知识库ID列表
                if knowledge_base_ids:
                    new_session.set_knowledge_base_ids(knowledge_base_ids)
                
                session.add(new_session)
                await session.commit()
                await session.refresh(new_session)
                logger.info(f"创建新会话: {session_id}, 知识库: {knowledge_base_ids}")
                return new_session

        except Exception as e:
            logger.exception(f"创建会话失败: {e}")
            raise

    async def update_session_knowledge_bases(self, session_id: str, knowledge_base_ids: List[str]) -> bool:
        """更新会话的知识库配置"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                chat_session = await session.get(ChatSession, session_id)
                if not chat_session:
                    # 如果会话不存在，创建一个
                    chat_session = ChatSession(
                        session_id=session_id,
                        created_at=datetime.now(),
                        last_activity=datetime.now(),
                        message_count=0
                    )
                    session.add(chat_session)

                # 更新知识库配置
                chat_session.set_knowledge_base_ids(knowledge_base_ids)
                chat_session.last_activity = datetime.now()

                await session.commit()
                logger.info(f"更新会话 {session_id} 的知识库配置: {knowledge_base_ids}")
                return True

        except Exception as e:
            logger.exception(f"更新会话知识库配置失败: {e}")
            return False

    async def update_session_activity(self, session_id: str, message_content: str = None,
        message_role: str = None) -> bool:
        """更新会话活动"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                chat_session = await session.get(ChatSession, session_id)
                if not chat_session:
                    # 如果会话不存在，创建一个
                    chat_session = ChatSession(
                        session_id=session_id,
                        created_at=datetime.now(),
                        last_activity=datetime.now(),
                        message_count=0
                    )
                    session.add(chat_session)

                # 更新活动时间和消息信息
                chat_session.last_activity = datetime.now()
                if message_content:
                    chat_session.last_message = message_content[:500]  # 限制长度
                    chat_session.message_count += 1
                if message_role:
                    chat_session.last_message_role = message_role

                await session.commit()
                return True

        except Exception as e:
            logger.exception(f"更新会话活动失败: {e}")
            return False

    async def get_session(self, session_id: str) -> Optional[ChatSession]:
        """获取单个会话"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                return await session.get(ChatSession, session_id)
        except Exception as e:
            logger.exception(f"获取会话失败: {e}")
            return None

    async def get_all_sessions(self, limit: int = 50, offset: int = 0) -> List[ChatSession]:
        """获取所有会话列表"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                from sqlalchemy import select
                stmt = (
                    select(ChatSession)
                    .order_by(ChatSession.last_activity.desc())
                    .limit(limit)
                    .offset(offset)
                )
                result = await session.execute(stmt)
                sessions = result.scalars().all()
                return list(sessions)

        except Exception as e:
            logger.exception(f"获取会话列表失败: {e}")
            return []

    async def delete_session(self, session_id: str) -> bool:
        """删除会话"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                chat_session = await session.get(ChatSession, session_id)
                if chat_session:
                    await session.delete(chat_session)
                    await session.commit()
                    logger.info(f"删除会话: {session_id}")
                    return True
                return False
        except Exception as e:
            logger.exception(f"删除会话失败: {e}")
            return False

    async def clear_old_sessions(self, days: int = 30) -> int:
        """清理旧会话"""
        try:
            engine = await self._get_engine()
            async with AsyncSession(engine) as session:
                from sqlalchemy import select, delete
                cutoff_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
                cutoff_date = cutoff_date.replace(day=cutoff_date.day - days)

                stmt = delete(ChatSession).where(ChatSession.last_activity < cutoff_date)
                result = await session.execute(stmt)
                await session.commit()

                deleted_count = result.rowcount()
                logger.info(f"清理了 {deleted_count} 个旧会话")
                return deleted_count

        except Exception as e:
            logger.exception(f"清理旧会话失败: {e}")
            return 0


# 全局会话管理器实例
session_manager = SessionManager()
